// Copyright 2011 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "shell/invoke_command.h"

#include <stdio.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>

#include <algorithm>

#include "base/logging.h"
#include "base/macros.h"

namespace flame {

static const int kStdin = 0;
static const int kStdout = 1;
static const int kStderr = 2;

static const uint8 kReadEnd = 0;
static const uint8 kWriteEnd = 1;

// Read from the file descriptor, appending to the string. Return false if we
// have reached the end of the file. Crash if there is an error.
static bool ReadAndAppendToString(
    int fd,
    string* s) {
  char buf[4];
  const ssize_t result = read(fd, buf, arraysize(buf));
  PCHECK(result >= 0);

  // Check for EOF.
  if (!result) {
    return false;
  }

  *s += string(buf, result);
  return true;
}

uint8 InvokeCommand(
    const string& command,
    string* stdout_contents,
    string* stderr_contents) {
  // Create pipes for the process's stdin, stdout, and stderr streams.
  int p_stdin[2];
  int p_stdout[2];
  int p_stderr[2];

  PCHECK(pipe(p_stdin) == 0);
  PCHECK(pipe(p_stdout) == 0);
  PCHECK(pipe(p_stderr) == 0);

  // Fork.
  const pid_t pid = fork();
  PCHECK(pid >= 0);

  // Is this the child process?
  if (!pid) {
    // Set up streams.
    PCHECK(dup2(p_stdin[kReadEnd], kStdin) == kStdin);
    PCHECK(dup2(p_stdout[kWriteEnd], kStdout) == kStdout);
    PCHECK(dup2(p_stderr[kWriteEnd], kStderr) == kStderr);

    // Close the unnecessary ends of the pipes.
    PCHECK(close(p_stdin[kWriteEnd]) == 0);
    PCHECK(close(p_stdout[kReadEnd]) == 0);
    PCHECK(close(p_stderr[kReadEnd]) == 0);

    // Call the shell.
    execl("/bin/sh", "sh", "-c", command.c_str(), NULL);
    perror("execl");
    exit(1);
  }

  // Close the unnecessary ends of the pipes.
  PCHECK(close(p_stdin[kReadEnd]) == 0);
  PCHECK(close(p_stdout[kWriteEnd]) == 0);
  PCHECK(close(p_stderr[kWriteEnd]) == 0);

  // We don't give the process any input.
  PCHECK(close(p_stdin[kWriteEnd]) == 0);

  // Keep reading until both streams are closed.
  for (;;) {
    const int stdout_read = p_stdout[kReadEnd];
    const int stderr_read = p_stderr[kReadEnd];

    // We want to read from stdout and stderr.
    fd_set read_fds;
    FD_ZERO(&read_fds);
    FD_SET(stdout_read, &read_fds);
    FD_SET(stderr_read, &read_fds);

    // Wait until one is ready for reading.
    const int nfds = std::max(stdout_read, stderr_read) + 1;
    PCHECK(select(nfds, &read_fds, NULL, NULL, NULL) > 0);

    // Figure out which stream to read from.
    const int ready_fd =
        FD_ISSET(stdout_read, &read_fds) ? stdout_read : stderr_read;
    string* contents =
        (ready_fd == stdout_read) ? stdout_contents : stderr_contents;

    // Read from the stream. If it has hit EOF, read from the other stream until
    // it's exhausted and then stop.
    if (!ReadAndAppendToString(ready_fd, contents)) {
      const int other_fd =
          (ready_fd == stdout_read) ? stderr_read : stdout_read;

      string* other_contents =
          (other_fd == stdout_read) ? stdout_contents : stderr_contents;

      while (ReadAndAppendToString(other_fd, other_contents)) {
      }

      break;
    }
  }

  // Wait until the process exits.
  int stat_loc;
  PCHECK(waitpid(pid, &stat_loc, 0) > 0);
  CHECK(WIFEXITED(stat_loc)) << "Process exited abnormally.";

  return WEXITSTATUS(stat_loc);
}

}  // namespace flame
